import time
from datetime import datetime

import pyrealsense2 as rs
import numpy as np
import socket
import json
import os
import cv2
from time import *
from Controller.WaterController import *

# 机械臂连接配置
robot_ip = "192.168.10.18"
robot_port = 8080
robot_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
robot_socket.settimeout(10)
robot_socket.connect((robot_ip, robot_port))
print(f"成功连接到机械臂 {robot_ip}:{robot_port}")


class RealSenseCamera:
    def __init__(self):
        self.pipeline = rs.pipeline()
        self.config = rs.config()
        self.config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)
        self.image_counter = 0
        self.save_dir = "dataset_retail/train/images"
        os.makedirs(self.save_dir, exist_ok=True)
        self.pipeline.start(self.config)
        print("相机初始化成功")

    def capture_photos(self, num_photos=1):
        """在当前位置拍摄指定数量的照片"""
        for _ in range(num_photos):
            frames = self.pipeline.wait_for_frames()
            color_frame = frames.get_color_frame()
            if not color_frame:
                continue

            color_img = np.asanyarray(color_frame.get_data())
            filename = f"image_{int(datetime.now().timestamp())}.jpg"
            save_path = os.path.join(self.save_dir, filename)
            cv2.imwrite(save_path, color_img)
            print(f"已保存图像 {filename}")

    def stop(self):
        self.pipeline.stop()
        print("相机已关闭")


def send_arm_command(command, max_retries=3):
    """
    发送json指令给机械臂动作，并根据device值处理响应，带重试机制
    :param command: json格式的机械臂命令
    :param max_retries: 最大重试次数
    :return: None
    """
    buffer = ""  # 用于累积接收到的数据
    first_message = None  # 存储第一条消息
    target_message = None  # 存储匹配目标device的消息
    retries = 0  # 重试计数

    if not robot_socket:
        print("未连接到机械臂")
        return None

    # 将字典转换为JSON字符串
    json_str = json.dumps(command)
    robot_socket.sendall((json_str + '\n').encode('utf-8'))
    print(f"已发送机械臂指令: {json_str}")

    while retries < max_retries:
        try:
            while True:
                try:
                    response = robot_socket.recv(1024).decode('utf-8')
                    if not response:
                        print("服务端断开连接")
                        break
                    buffer += response

                    while '\n' in buffer:
                        message, buffer = buffer.split('\n', 1)
                        try:
                            parsed_message = json.loads(message)
                            print(f"Received message: {parsed_message}")

                            # 如果是第一条消息，存储用于state检查
                            if first_message is None:
                                first_message = parsed_message

                            # 检查device值
                            device_value = parsed_message.get('device')
                            if device_value == 3:
                                print(f"检测到device=3，跳出循环")
                                return  # 直接返回

                            # 如果device为0或1，且还未找到目标消息
                            if device_value in (0, 1) and target_message is None:
                                target_message = parsed_message
                                state_ok = (first_message.get('state', False) == True or
                                            first_message.get('receive_state', False) == True)
                                trajectory_state_ok = target_message.get('trajectory_state', False) == True
                                print(f"State (from first message) OK: {state_ok}, "
                                      f"Trajectory state (from target message) OK: {trajectory_state_ok}")
                                if state_ok and trajectory_state_ok:
                                    print(f"指令执行完成，device={device_value}，状态验证通过")
                                    return
                                else:
                                    print("指令执行未完成或状态异常:")
                                    print(f"First message: {first_message}")
                                    print(f"Target message: {target_message}")
                                    raise RuntimeError("机械臂指令执行失败，状态未达到预期")

                        except json.JSONDecodeError:
                            buffer = message + '\n' + buffer
                            break
                except socket.timeout:
                    print(f"接收响应超时（第 {retries + 1} 次尝试），可能是机械臂未正确响应")
                    if first_message and target_message:
                        state_ok = (first_message.get('state', False) == True or
                                    first_message.get('receive_state', False) == True)
                        trajectory_state_ok = target_message.get('trajectory_state', False) == True
                        if state_ok and trajectory_state_ok:
                            print(f"指令执行完成，超时前已收到有效消息，状态验证通过")
                            return
                    retries += 1
                    if retries < max_retries:
                        print(f"重试 {retries + 1}/{max_retries}")
                        buffer = ""  # 清空缓冲区，重试
                        robot_socket.sendall((json_str + '\n').encode('utf-8'))  # 重新发送指令
                        sleep(1)  # 短暂等待后重试
                    break

        except socket.timeout:
            continue  # 外层循环处理重试

    raise RuntimeError(f"机械臂响应超时，经过 {max_retries} 次重试仍未成功")


def water_to_target(target_point, controller: Water):
    if ',' in target_point:
        split = target_point.split(",")
        task = controller.move_to_pose(split[1], split[2], split[0], 3)
    else:
        task = controller.move_to_marker(target_point)
    if task["status"] == "OK":
        try:
            # publisher.updateStates(woosh_state=WooshModel.RobotState.kTask)
            while True:
                sleep(1)
                state = controller.get_move_status()
                if state["move_status"] != "running":
                    # publisher.updateStates(woosh_state=WooshModel.RobotState.kfree)
                    controller.cancel_move()
                    break
            return True
        except:
            return False
    else:
        return False


def main():
    # 初始化相机
    camera = RealSenseCamera()
    sleep(1)
    camera.capture_photos(num_photos=2)
    init_joint = {"command": "movej", "joint": [1957, 117059, -89232, -846, 114787, -7614], "v": 30, "r": 0,
                  "trajectory_connect": 0}
    Gripper_realease = {"command": "set_gripper_release", "speed": 500, "force": 200, "block": True}
    photo_guodudian1 = {"command": "movej", "joint": [1237, 67912, -129954, 6192, 63572, -2760], "v": 30, "r": 0,
                        "trajectory_connect": 0}
    # 拍照前的过渡点2
    photo_guodudian2 = {"command": "movej", "joint": [-91958, 62804, -132521, 1911, 70301, 90531], "v": 30, "r": 0,
                        "trajectory_connect": 0}
    photo1_command = {"command": "movej", "joint": [-89772, 119139, -117335, 1933, -3617, 86481], "v": 30, "r": 0,
                      "trajectory_connect": 0}
    photo2_command = {"command": "movel", "pose": [624603, 326323, 259, 1665, -1546, 1443], "v": 30, "r": 0,
                      "trajectory_connect": 0}
    photo3_command = {"command": "movel", "pose": [625530, 433062, 226, 1667, -1546, 1436], "v": 30, "r": 0,
                      "trajectory_connect": 0}
    point_dict = ["photo1", "photo2", "photo3"]

    # 让机器人走3个点位
    for robot_point in point_dict:
        # 机器人导航到robot_point
        water = Water("http://192.168.10.10:9001", map_code=None)
        water_to_target(robot_point, water)
        send_arm_command(photo_guodudian1)
        sleep(1)
        send_arm_command(photo_guodudian2)
        sleep(1)
        # 让升降模组走3个层高
        hight_offset = 1050
        for layer in range(3):
            init_hight = 270
            # 升降模组走到绝对位置
            hight = init_hight + layer * hight_offset
            # 升降模组运动
            # 升降命令
            lift_up = {"command": "set_lift_height", "height": hight, "speed": 80}
            send_arm_command(lift_up)
            # 机械臂运动到第一个点moveJ
            send_arm_command(photo1_command)
            sleep(2)
            # 拍照存储
            camera.capture_photos()
            # 机械臂运动到第二个点moveL
            send_arm_command(photo2_command)
            sleep(2)
            # 拍照存储
            camera.capture_photos()
            # 机械臂运动到第三个点moveL
            send_arm_command(photo3_command)
            sleep(2)
            # 拍照存储
            camera.capture_photos()
            # 回到初始点
            send_arm_command(photo1_command)
        send_arm_command(init_joint)


if __name__ == "__main__":
    main()
